import gc
import os
import random
import zipfile
from io import BytesIO

import colorsys
import cv2
import numpy as np
import rasterio
from rasterio import CRS
from rasterio.io import MemoryFile
from rasterio.warp import calculate_default_transform, reproject, Resampling
import matplotlib.pyplot as plt
from PIL import Image, ImageColor


def read_land_cover_2():
    parent_directory = 'E:\\map data'
    file = '2020LC030.zip'
    keys_and_color0 = {
        10: [254, 246, 201], #耕地
        20: [16, 118, 75], #森林
        30: [172, 212, 90], #草地
        40: [57, 119, 115], #灌木地
        50: [124, 209, 245], #湿地
        60: [0, 87, 155], #水体
        70: [150, 200, 100], #苔原
        80: [147, 45, 16], #人造地表
        90: [206, 203, 206], #裸地
        100: [214, 242, 255], #冰川
    }

    keys_and_color = {}

    all_width = 21600; all_height = 10800
    all_image = np.zeros((all_height, all_width, 4), dtype=np.uint8) #RGBA

    local_width = -1; local_height = -1
    lon_size = 360; lat_size = 180

    arr0 = np.array([[2, 3, 1], [2, 9, 12]])
    arr1 = np.array([[2, 2, 9], [12, 22, 33]])

    directory_name = os.path.join(parent_directory, file)
    zip_file_info = zipfile.ZipFile(directory_name, 'r')
    file_list = zip_file_info.namelist()

    file_list_filtered = [s for s in file_list if s.endswith('.tif')]
    f_count = 0

    for file in file_list:
        if file.endswith('.zip'):
            print(f'{file} {f_count}/{len(file_list_filtered)}')

            single_zip_file = zip_file_info.open(file)
            zip_content_info = single_zip_file.read()

            zip_content = zipfile.ZipFile(BytesIO(zip_content_info), 'r')
            local_file_list = zip_content.namelist()

            for single_file_name in local_file_list:
                if single_file_name.endswith('.tif'):
                    single_file = zip_content.open(single_file_name)
                    content = single_file.read()
                    with (MemoryFile(content) as memfile):
                        # 目标投影：WGS84经纬度（EPSG:4326）
                        dst_crs = CRS.from_epsg(4326)

                        # 打开内存文件中的 TIFF 数据集
                        with memfile.open() as utm_src:
                            # 示例：读取影像数据（numpy 数组）
                            image_nparray_0 = utm_src.read() # shape: (波段数, 高度, 宽度)
                            image_nparray = image_nparray_0[0]

                            crs = utm_src.crs
                            transform, width, height = calculate_default_transform(utm_src.crs, dst_crs, utm_src.width,
                                                                                   utm_src.height,
                                                                                   *utm_src.bounds)

                            dst_meta = utm_src.meta.copy()
                            dst_meta.update({
                                "crs": dst_crs, "transform": transform, "width": width, "height": height,
                            })

                            # 创建内存文件存储经纬度栅格
                            with MemoryFile() as wgs84_mem:
                                with wgs84_mem.open(**dst_meta) as wgs84_dst:
                                    # 执行投影转换（重采样）
                                    reproject(
                                        source=rasterio.band(utm_src, 1),  # 源波段
                                        destination=rasterio.band(wgs84_dst, 1),  # 目标波段
                                        src_transform=utm_src.transform,
                                        src_crs=utm_src.crs,
                                        dst_transform=transform,
                                        dst_crs=dst_crs,
                                        resampling=Resampling.nearest  # 插值方法：连续数据用双线性，分类数据用nearest
                                    )

                                    # ---------------------- 验证转换结果 ----------------------
                                    print("=== UTM转经纬度结果验证 ===")
                                    print(f"源UTM投影: {utm_src.crs.to_epsg()} (UTM 50N)")
                                    print(f"目标经纬度投影: {wgs84_dst.crs.to_epsg()} (WGS84)")
                                    print(f"源UTM尺寸: {utm_src.width}x{utm_src.height} (米单位)")
                                    print(f"目标经纬度尺寸: {wgs84_dst.width}x{wgs84_dst.height} (度单位)")
                                    print(f"经纬度变换矩阵: {wgs84_dst.transform}")  # 分辨率为度（如0.00009°）
                                    print(f"经纬度栅格边界: {wgs84_dst.bounds}")  # 经纬度范围

                                    current_bounds = wgs84_dst.bounds
                                    current_left = current_bounds.left; current_top = current_bounds.top
                                    current_right = current_bounds.right; current_bottom = current_bounds.bottom

                                    lon_gap = current_right - current_left; lat_gap = current_bottom - current_top
                                    ratio0 = lon_gap / lat_gap

                                    pixel_width = round(all_width * lon_gap / lon_size)
                                    pixel_height = round(all_height * abs(lat_gap) / lat_size)

                                    start_x = round((current_left - (-180)) / lon_size * all_width)
                                    start_y = round((1 - (current_top - (-90)) / lat_size) * all_height)

                                    shape_ratio0 = wgs84_dst.width / wgs84_dst.height
                                    # 读取转换后的数据
                                    wgs84_data = wgs84_dst.read(1)
                                    wgs84_data[wgs84_data == 255] = 0

                                    wgs84_data_rgba_ = np.zeros((wgs84_dst.height, wgs84_dst.width, 4), dtype=np.uint8)

                                    pixel_unique = np.unique(wgs84_data)
                                    for pixel in pixel_unique:
                                        if pixel != 0:
                                            if pixel not in keys_and_color:
                                                h = random.random(); s = 0.8; v = 0.5
                                                color0 = colorsys.hsv_to_rgb(h, s, v)
                                                color_r = int(color0[0] * 255); color_g = int(color0[1] * 255)
                                                color_b = int(color0[2] * 255)

                                                if pixel in keys_and_color0:
                                                    local_color = keys_and_color0[pixel]
                                                    color_r = local_color[0]; color_g = local_color[1]; color_b = local_color[2]
                                                keys_and_color[pixel] = np.array([color_r, color_g, color_b], dtype=np.uint8)

                                            single_color = keys_and_color[pixel]
                                            wgs84_data_rgba_[:, :, 0][wgs84_data == pixel] = single_color[0]
                                            wgs84_data_rgba_[:, :, 1][wgs84_data == pixel] = single_color[1]
                                            wgs84_data_rgba_[:, :, 2][wgs84_data == pixel] = single_color[2]
                                            wgs84_data_rgba_[:, :, 3][wgs84_data == pixel] = 255



                                    # wgs84_data_rgba = np.stack([wgs84_data_r, wgs84_data_g, wgs84_data_b, wgs84_data_a], dtype=np.uint8)
                                    wgs84_data1 = cv2.resize(wgs84_data_rgba_, (pixel_width, pixel_height),
                                                     interpolation=cv2.INTER_AREA)

                                    image_area_to_fill = all_image[start_y:start_y + pixel_height, start_x:start_x + pixel_width]
                                    condition = image_area_to_fill[:, :, 3] < wgs84_data1[:, :, 3]

                                    image_area_to_fill[condition] = wgs84_data1[condition]

                                    del wgs84_data_rgba_
                                    gc.collect()
                                    # all_image[start_y:start_y + pixel_height, start_x:start_x + pixel_width] = np.maximum(all_image[start_y:start_y + pixel_height, start_x:start_x + pixel_width], wgs84_data1)
                                    t = 11
            f_count += 1

    image = Image.fromarray(all_image)
    image.show()
    return file_list

if __name__ == '__main__':
    read_land_cover_2()